package org.dlx.drpc.core.client;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

//一个简单的连接池管理
public class RPCPool
{
    private RPCClient client;

    private HashMap<RPCConn, Boolean> pool = new HashMap<>();

    private AtomicBoolean close = new AtomicBoolean();

    public RPCPool(RPCClient client) {
        this.client = client;
        this.close.set(false);
    }

    public <T> T send(Object payload, Class<T> cla, long ms) throws Exception {
        RPCConn conn = load(false);
        try {
            T re = this.client.send(conn, payload, cla, ms);
            this.store(conn);//链接用完 就恢复
            return re;
        } catch (IOException io) {
            this.remove(conn);//移除旧链接
            conn = load(true);
            try {
                T re = this.client.send(conn, payload, cla, ms);
                this.store(conn);
                return re;
            } catch (Exception e) {
                remove(conn);
                throw e;
            }
        }catch (Exception e) {
            this.remove(conn);//发生IO以外的异常 移除掉CONN
            throw e;
        }
    }

    public void close() {
        synchronized (pool) {
            close.set(true);
            for (Map.Entry<RPCConn, Boolean> entry : pool.entrySet())
            {
                entry.getKey().close();
            }
            pool.clear();
            client.close();
        }
    }

    //从连接池获取一个连接
    private RPCConn load(boolean forceCreate) throws Exception{
        synchronized (pool) {
            if(close.get()) {
                throw new Exception("连接池已经关闭");
            }
            if (!forceCreate) {
                for (Map.Entry<RPCConn, Boolean> entry : pool.entrySet()) {
                    if (!entry.getValue()) {
                        entry.setValue(true);
                        return entry.getKey();
                    }
                }
            }
            //没有找到或者强制创建新的链接
            RPCConn cn = client.newConn();
            pool.put(cn, true);
            return cn;
        }
    }

    //释放连接占用
    private void store(RPCConn conn) {
        synchronized (pool) {
            if(close.get())
                return;
            pool.put(conn, false);
        }
    }

    private void remove(RPCConn conn) {
        synchronized (pool) {
            if(close.get())//调用过close后pool已经空了
                return;
            pool.remove(conn);
        }
    }
}
